package com.cetc.sdp.kmga.cs.util;

import org.apache.commons.logging.Log;
import org.apache.commons.logging.LogFactory;
import org.apache.hadoop.conf.Configuration;
import org.apache.hadoop.hbase.HBaseConfiguration;
import org.apache.hadoop.hbase.TableName;
import org.apache.hadoop.hbase.client.*;
import org.apache.hadoop.hbase.exceptions.ConnectionClosingException;
import org.apache.hadoop.hbase.io.ImmutableBytesWritable;
import org.apache.hadoop.hbase.mapreduce.TableOutputCommitter;
import org.apache.hadoop.hbase.util.Bytes;
import org.apache.hadoop.mapreduce.*;

import java.io.IOException;
import java.time.LocalDateTime;
import java.util.HashMap;
import java.util.Map;
import java.util.Queue;
import java.util.concurrent.LinkedBlockingQueue;
import java.util.regex.Pattern;

/**
 * Hbase Output 连接池
 *
 * @description:
 * @author： DengQiang
 * @date: 2017/10/25 15:38
 */
public class PooledMultiTableOutputFormat extends OutputFormat<ImmutableBytesWritable, Mutation> {
    /**
     * Set this to {@link #WAL_OFF} to turn off write-ahead logging (WAL)
     */
    public static final String WAL_PROPERTY = "hbase.mapreduce.multitableoutputformat.wal";
    /**
     * Property value to use write-ahead logging
     */
    public static final boolean WAL_ON = true;
    /**
     * Property value to disable write-ahead logging
     */
    public static final boolean WAL_OFF = false;

    /**
     * Record writer for outputting to multiple HTables.
     */
    protected static class MultiTableRecordWriter extends
            RecordWriter<ImmutableBytesWritable, Mutation> {
        private static final Log LOG = LogFactory.getLog(PooledMultiTableOutputFormat.class);
        Connection connection;
        Map<ImmutableBytesWritable, BufferedMutator> mutatorMap = new HashMap<>();
        Configuration conf;
        boolean useWriteAheadLogging;
        private volatile boolean connectionExp = false;

        /**
         * @param conf                 HBaseConfiguration to used
         * @param useWriteAheadLogging whether to use write ahead logging. This can be turned off (
         *                             <tt>false</tt>) to improve performance when bulk loading data.
         */
        public MultiTableRecordWriter(Configuration conf,
                                      boolean useWriteAheadLogging) {
            LOG.debug("Created new MultiTableRecordReader with WAL "
                    + (useWriteAheadLogging ? "on" : "off"));
            this.conf = conf;
            this.useWriteAheadLogging = useWriteAheadLogging;
        }

        /**
         * @param tableName the name of the table, as a string
         * @return the named mutator
         * @throws IOException if there is a problem opening a table
         */
        BufferedMutator getBufferedMutator(ImmutableBytesWritable tableName) throws IOException {
            if (this.connection == null) {
                this.connection = ConnectionFactory.createConnection(conf);
            }
            if (!mutatorMap.containsKey(tableName)) {
                String name = Bytes.toString(tableName.get());
                LOG.debug("Opening HTable \"" + name + "\" for writing");
                System.out.println(LocalDateTime.now() + " Opening HTable \"" + name + "\" for writing");
                TableName tn = TableName.valueOf(tableName.get());
                BufferedMutator mutator = connection.getBufferedMutator(tn);
                mutatorMap.put(tableName, mutator);
//                BufferedMutatorParams params = new BufferedMutatorParams(tn);
//                params.listener(new BufferedMutator.ExceptionListener() {
//                    @Override
//                    public void onException(RetriesExhaustedWithDetailsException exception, BufferedMutator mutator) throws RetriesExhaustedWithDetailsException {
//                        connectionExp = true;
//                        throw exception;
//                    }
//                });

//                /**
//                 * 由于部分数据每个月会保存的对应的月表中，因此当月表变更时，需要关闭与上个月的表的不用连接
//                 * 这种月表的表名中包含年月信息，模式是: .+(\d{6}).*
//                 * 首先通过表名识别出这样的表，然后判断是否有上个月的表需要关闭
//                 */
//                Matcher matcher = monthTablePattern.matcher(name);
//                if (matcher.find()) {
//                    try {
//                        String ym = matcher.group(1);
//                        DateTimeFormatter formatter = DateTimeFormatter.ofPattern("yyyyMM");
//                        TemporalAccessor accessor = formatter.parse(ym);
//                        LocalDate ld = LocalDate.of(accessor.get(ChronoField.YEAR), accessor.get(ChronoField.MONTH_OF_YEAR), 1);
//                        LocalDate be = ld.minusMonths(1);
//                        String oldTableName = name.replace(ym, formatter.format(be));
//                        ImmutableBytesWritable ibw = new ImmutableBytesWritable(Bytes.toBytes(oldTableName));
//                        BufferedMutator oldMutator = mutatorMap.remove(ibw);
//                        if (oldMutator != null) {
//                            System.out.println("释放与上个月表的连接，Table: " + oldTableName);
//                            oldMutator.close();
//                        }
//                    } catch (Exception e) {
//                        e.printStackTrace();
//                    }
//                }
            }
            return mutatorMap.get(tableName);
        }

        @Override
        public void close(TaskAttemptContext context) throws IOException {
            if (connectionExp) {
                this.connection = null;
                throw new ConnectionClosingException("HBase连接异常.");
            }
            try {
                for (BufferedMutator mutator : mutatorMap.values()) {
                    mutator.close();
                }
                mutatorMap.clear();
                writerPool.offer(this);
            } catch (IOException e) {
                System.out.println(e.getMessage());
                if (this.connection != null) {
                    this.connection.close();
                    mutatorMap.clear();
                }
                this.connection = null;
                throw e;
            }
        }

        /**
         * Writes an action (Put or Delete) to the specified table.
         *
         * @param tableName the table being updated.
         * @param action    the update, either a put or a delete.
         * @throws IllegalArgumentException if the action is not a put or a delete.
         */
        @Override
        public void write(ImmutableBytesWritable tableName, Mutation action) throws IOException {
            try {
                BufferedMutator mutator = getBufferedMutator(tableName);
                // The actions are not immutable, so we defensively copy them
                if (action instanceof Put) {
                    Put put = (Put) action;
                    put.setDurability(useWriteAheadLogging ? Durability.ASYNC_WAL
                            : Durability.SKIP_WAL);
                    mutator.mutate(action);
                } else if (action instanceof Delete) {
//                Delete delete = new Delete((Delete) action);
                    mutator.mutate(action);
                } else {
                    throw new IllegalArgumentException(
                            "action must be either Delete or Put");
                }
            } catch (IOException e) {
                e.printStackTrace();
                this.connectionExp = true;
                throw e;
            }
        }
    }

    @Override
    public void checkOutputSpecs(JobContext context) throws IOException,
            InterruptedException {
        // we can't know ahead of time if it's going to blow up when the user
        // passes a table name that doesn't exist, so nothing useful here.
    }

    @Override
    public OutputCommitter getOutputCommitter(TaskAttemptContext context)
            throws IOException, InterruptedException {
        return new TableOutputCommitter();
    }

    @Override
    public RecordWriter<ImmutableBytesWritable, Mutation> getRecordWriter(TaskAttemptContext context)
            throws IOException, InterruptedException {
        PooledMultiTableOutputFormat.MultiTableRecordWriter writer = writerPool.poll();
        if (writer == null) {
            System.out.println(LocalDateTime.now() + " Get new PooledMultiTableOutputFormat.MultiTableRecordWriter.");
            Configuration conf = context.getConfiguration();
            writer = new PooledMultiTableOutputFormat.MultiTableRecordWriter(HBaseConfiguration.create(conf),
                    conf.getBoolean(WAL_PROPERTY, WAL_ON));
        }

        if ((writer.connection != null && (writer.connectionExp || writer.connection.isClosed()))) {
            writer.mutatorMap.clear();
            Configuration conf = context.getConfiguration();
            writer = new PooledMultiTableOutputFormat.MultiTableRecordWriter(HBaseConfiguration.create(conf),
                    conf.getBoolean(WAL_PROPERTY, WAL_ON));
        }

        return writer;
    }

    private static Queue<MultiTableRecordWriter> writerPool = new LinkedBlockingQueue<>();

    private static Pattern monthTablePattern = Pattern.compile(".+(\\d{6}).*");

}
